{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp data.transforms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Time Series Data Augmentation\n",
    "\n",
    "> Functions used to transform TSTensors (Data Augmentation)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from tsai.imports import *\n",
    "from fastai.vision.augment import RandTransform\n",
    "from tsai.utils import *\n",
    "from tsai.data.external import *\n",
    "from tsai.data.core import *\n",
    "from tsai.data.preprocessing import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from scipy.interpolate import CubicSpline\n",
    "from scipy.ndimage import convolve1d, zoom\n",
    "import pywt\n",
    "from pyts.image.gaf import GramianAngularField"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dsid = 'NATOPS'\n",
    "X, y, splits = get_UCR_data(dsid, return_split=False)\n",
    "tfms = [None, Categorize()]\n",
    "batch_tfms = TSStandardize()\n",
    "dls = get_ts_dls(X, y, tfms=tfms, splits=splits, batch_tfms=batch_tfms, bs=128)\n",
    "xb, yb = next(iter(dls.train))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSIdentity(RandTransform):\n",
    "    \"Applies the identity tfm to a `TSTensor` batch\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=None, **kwargs): \n",
    "        self.magnitude = magnitude \n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor): return o"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSIdentity()(xb, split_idx=0).shape, xb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "# partial(TSShuffle_HLs, ex=0), \n",
    "class TSShuffle_HLs(RandTransform):\n",
    "    \"Randomly shuffles HIs/LOs of an OHLC `TSTensor` batch\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=1., ex=None, **kwargs): \n",
    "        self.magnitude, self.ex = magnitude, ex\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        timesteps = o.shape[-1] // 4\n",
    "        pos_rand_list = np.random.choice(np.arange(timesteps),size=random.randint(1, timesteps),replace=False)\n",
    "        rand_list = pos_rand_list * 4\n",
    "        highs = rand_list + 1\n",
    "        lows = highs + 1\n",
    "        a = np.vstack([highs, lows]).flatten('F')\n",
    "        b = np.vstack([lows, highs]).flatten('F')\n",
    "        output = o.clone()\n",
    "        output[...,a] = output[...,b]\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSShuffle_HLs()(xb, split_idx=0).shape, xb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "# partial(TSShuffleSteps, ex=0), \n",
    "class TSShuffleSteps(RandTransform):\n",
    "    \"Randomly shuffles consecutive sequence datapoints in batch\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=1., ex=None, **kwargs): \n",
    "        self.magnitude, self.ex = magnitude, ex\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        odd = 1 - o.shape[-1]%2\n",
    "        r = np.random.randint(2)\n",
    "        timesteps = o.shape[-1] // 2\n",
    "        pos_rand_list = np.random.choice(np.arange(0, timesteps - r * odd), size=random.randint(1, timesteps - r * odd),replace=False) * 2 + 1 + r\n",
    "        a = np.vstack([pos_rand_list, pos_rand_list - 1]).flatten('F')\n",
    "        b = np.vstack([pos_rand_list - 1, pos_rand_list]).flatten('F')\n",
    "        output = o.clone()\n",
    "        output[...,a] = output[...,b]\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXcAAAD4CAYAAAAXUaZHAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/MnkTPAAAACXBIWXMAAAsTAAALEwEAmpwYAAAPSklEQVR4nO3cf6yeZX3H8fdnFPyBCwU5a1jb7JDYaIiJQk4Qx2Ic3RZ+GMsfyDCbdqxL/0GH00Sr/5gl+wOTRcVkYWmoWjYmEtS0UXSSgjH+Qefhx5AfEs4Y2HYtPSpUHXGO+d0f5+p2KC3nOT3Pcx569f1KntzXfV3X/dzfO20+5z7XuZ8nVYUkqS+/Me4CJEnDZ7hLUocMd0nqkOEuSR0y3CWpQyvGXQDA2WefXZOTk+MuQ5JOKPfdd9+Pq2riaGOviHCfnJxkenp63GVI0gklydPHGnNZRpI6ZLhLUocMd0nqkOEuSR0y3CWpQ4a7JHXIcJekDhnuktQhw12SOvSK+ISqTm6TW74xsvd+6oYrRvbeg+r9+vTK5J27JHVooDv3JCuBm4E3AwX8OfA48GVgEngKuLqqnk0S4EbgcuB54M+q6v5hFy4dL++kdTIYdFnmRuBbVXVVktOA1wKfAHZV1Q1JtgBbgI8BlwHr2uttwE1tqxPEqMLP4Fs+/htqwWWZJGcA7wC2AVTVr6rqOWADsL1N2w5c2dobgFtqzr3AyiTnDLluSdLLGGTN/VxgFvhCkgeS3JzkdGBVVe1vcw4Aq1p7NbBn3vF7W9+LJNmcZDrJ9Ozs7PFfgSTpJQZZllkBXAB8sKp2J7mRuSWY/1NVlaQWc+Kq2gpsBZiamlrUsePkeq30Ui4DvfIMcue+F9hbVbvb/h3Mhf0zh5db2vZgG98HrJ13/JrWJ0laJgveuVfVgSR7kryxqh4H1gOPttdG4Ia23dEO2Ql8IMltzP0h9dC85RtJWjJ/U1jYoE/LfBC4tT0p8yRwLXN3/bcn2QQ8DVzd5t7J3GOQM8w9CnntUCuWJC1ooHCvqgeBqaMMrT/K3AKuW1pZms+7FEmL5SdUJalDhrskdchwl6QOGe6S1CHDXZI6ZLhLUocMd0nqkOEuSR0y3CWpQ4a7JHXIcJekDhnuktQhw12SOmS4S1KHDHdJ6pDhLkkdMtwlqUOGuyR1yHCXpA4Z7pLUIcNdkjpkuEtShwx3SeqQ4S5JHTLcJalDA4V7kqeS/CDJg0mmW99ZSe5K8kTbntn6k+RzSWaSPJTkglFegCTppRZz5/77VfXWqppq+1uAXVW1DtjV9gEuA9a112bgpmEVK0kazFKWZTYA21t7O3DlvP5bas69wMok5yzhPJKkRRo03Av4dpL7kmxufauqan9rHwBWtfZqYM+8Y/e2vhdJsjnJdJLp2dnZ4yhdknQsKwac93tVtS/JbwF3Jfnh/MGqqiS1mBNX1VZgK8DU1NSijpUkvbyB7tyral/bHgS+BlwIPHN4uaVtD7bp+4C18w5f0/okSctkwXBPcnqS3zzcBv4IeBjYCWxs0zYCO1p7J/D+9tTMRcChecs3kqRlMMiyzCrga0kOz/+nqvpWku8DtyfZBDwNXN3m3wlcDswAzwPXDr1qSdLLWjDcq+pJ4C1H6f8JsP4o/QVcN5TqJEnHxU+oSlKHDHdJ6pDhLkkdMtwlqUOGuyR1yHCXpA4Z7pLUIcNdkjpkuEtShwx3SeqQ4S5JHTLcJalDhrskdchwl6QOGe6S1CHDXZI6ZLhLUocMd0nqkOEuSR0y3CWpQ4a7JHXIcJekDhnuktQhw12SOjRwuCc5JckDSb7e9s9NsjvJTJIvJzmt9b+q7c+08ckR1S5JOobF3LlfDzw2b/9TwGeq6g3As8Cm1r8JeLb1f6bNkyQto4HCPcka4Arg5rYf4BLgjjZlO3Bla29o+7Tx9W2+JGmZDHrn/lngo8Cv2/7rgeeq6oW2vxdY3dqrgT0AbfxQmy9JWiYLhnuSdwEHq+q+YZ44yeYk00mmZ2dnh/nWknTSG+TO/WLg3UmeAm5jbjnmRmBlkhVtzhpgX2vvA9YCtPEzgJ8c+aZVtbWqpqpqamJiYkkXIUl6sQXDvao+XlVrqmoSuAa4u6r+BLgHuKpN2wjsaO2dbZ82fndV1VCrliS9rKU85/4x4MNJZphbU9/W+rcBr2/9Hwa2LK1ESdJirVh4yv+rqu8A32ntJ4ELjzLnl8B7hlCbJOk4+QlVSeqQ4S5JHTLcJalDhrskdchwl6QOGe6S1CHDXZI6ZLhLUocMd0nqkOEuSR0y3CWpQ4a7JHXIcJekDhnuktQhw12SOmS4S1KHDHdJ6pDhLkkdMtwlqUOGuyR1yHCXpA4Z7pLUIcNdkjpkuEtShwx3SerQguGe5NVJ/iXJvyZ5JMlft/5zk+xOMpPky0lOa/2vavszbXxyxNcgSTrCigHm/BdwSVX9IsmpwPeSfBP4MPCZqrotyd8Dm4Cb2vbZqnpDkmuATwF/PKL6mdzyjVG9NU/dcMXI3luSRmnBO/ea84u2e2p7FXAJcEfr3w5c2dob2j5tfH2SDKtgSdLCBlpzT3JKkgeBg8BdwL8Bz1XVC23KXmB1a68G9gC08UPA64/ynpuTTCeZnp2dXdJFSJJebKBwr6r/qaq3AmuAC4E3LfXEVbW1qqaqampiYmKpbydJmmdRT8tU1XPAPcDbgZVJDq/ZrwH2tfY+YC1AGz8D+MkwipUkDWaQp2Umkqxs7dcAfwg8xlzIX9WmbQR2tPbOtk8bv7uqaog1S5IWMMjTMucA25OcwtwPg9ur6utJHgVuS/I3wAPAtjZ/G/APSWaAnwLXjKBuSVo2J+JTeQuGe1U9BJx/lP4nmVt/P7L/l8B7hlKdJOm4+AlVSeqQ4S5JHTLcJalDhrskdchwl6QOGe6S1CHDXZI6ZLhLUocMd0nqkOEuSR0y3CWpQ4a7JHXIcJekDhnuktQhw12SOmS4S1KHDHdJ6pDhLkkdMtwlqUOGuyR1yHCXpA4Z7pLUIcNdkjpkuEtShwx3SerQguGeZG2Se5I8muSRJNe3/rOS3JXkibY9s/UnyeeSzCR5KMkFo74ISdKLDXLn/gLwkao6D7gIuC7JecAWYFdVrQN2tX2Ay4B17bUZuGnoVUuSXtaC4V5V+6vq/tb+OfAYsBrYAGxv07YDV7b2BuCWmnMvsDLJOcMuXJJ0bItac08yCZwP7AZWVdX+NnQAWNXaq4E98w7b2/qOfK/NSaaTTM/Ozi62bknSyxg43JO8DvgK8KGq+tn8saoqoBZz4qraWlVTVTU1MTGxmEMlSQsYKNyTnMpcsN9aVV9t3c8cXm5p24Otfx+wdt7ha1qfJGmZDPK0TIBtwGNV9el5QzuBja29Edgxr//97amZi4BD85ZvJEnLYMUAcy4G3gf8IMmDre8TwA3A7Uk2AU8DV7exO4HLgRngeeDaYRYsSVrYguFeVd8Dcozh9UeZX8B1S6xLkrQEfkJVkjpkuEtShwx3SeqQ4S5JHTLcJalDhrskdchwl6QOGe6S1CHDXZI6ZLhLUocMd0nqkOEuSR0y3CWpQ4a7JHXIcJekDhnuktQhw12SOmS4S1KHDHdJ6pDhLkkdMtwlqUOGuyR1yHCXpA4Z7pLUoQXDPcnnkxxM8vC8vrOS3JXkibY9s/UnyeeSzCR5KMkFoyxeknR0g9y5fxG49Ii+LcCuqloH7Gr7AJcB69prM3DTcMqUJC3GguFeVd8FfnpE9wZge2tvB66c139LzbkXWJnknCHVKkka0PGuua+qqv2tfQBY1dqrgT3z5u1tfS+RZHOS6STTs7Ozx1mGJOlolvwH1aoqoI7juK1VNVVVUxMTE0stQ5I0z/GG+zOHl1va9mDr3wesnTdvTeuTJC2j4w33ncDG1t4I7JjX//721MxFwKF5yzeSpGWyYqEJSb4EvBM4O8le4JPADcDtSTYBTwNXt+l3ApcDM8DzwLUjqFmStIAFw72q3nuMofVHmVvAdUstSpK0NH5CVZI6ZLhLUocMd0nqkOEuSR0y3CWpQ4a7JHXIcJekDhnuktQhw12SOmS4S1KHDHdJ6pDhLkkdMtwlqUOGuyR1yHCXpA4Z7pLUIcNdkjpkuEtShwx3SeqQ4S5JHTLcJalDhrskdchwl6QOGe6S1CHDXZI6NJJwT3JpkseTzCTZMopzSJKObejhnuQU4O+Ay4DzgPcmOW/Y55EkHdso7twvBGaq6smq+hVwG7BhBOeRJB1Dqmq4b5hcBVxaVX/R9t8HvK2qPnDEvM3A5rb7RuDxoRZybGcDP16mc42D13fi6/0avb7h+Z2qmjjawIplKuAlqmorsHW5z5tkuqqmlvu8y8XrO/H1fo1e3/IYxbLMPmDtvP01rU+StExGEe7fB9YlOTfJacA1wM4RnEeSdAxDX5apqheSfAD4Z+AU4PNV9ciwz7MEy74UtMy8vhNf79fo9S2Dof9BVZI0fn5CVZI6ZLhLUodOmnDv+SsRkqxNck+SR5M8kuT6cdc0KklOSfJAkq+Pu5ZhS7IyyR1JfpjksSRvH3dNw5bkr9r/0YeTfCnJq8dd01Ik+XySg0kentd3VpK7kjzRtmeOo7aTItxPgq9EeAH4SFWdB1wEXNfZ9c13PfDYuIsYkRuBb1XVm4C30Nl1JlkN/CUwVVVvZu6Bi2vGW9WSfRG49Ii+LcCuqloH7Gr7y+6kCHc6/0qEqtpfVfe39s+ZC4XV461q+JKsAa4Abh53LcOW5AzgHcA2gKr6VVU9N9aiRmMF8JokK4DXAv8x5nqWpKq+C/z0iO4NwPbW3g5cuZw1HXayhPtqYM+8/b10GH4ASSaB84HdYy5lFD4LfBT49ZjrGIVzgVngC23Z6eYkp4+7qGGqqn3A3wI/AvYDh6rq2+OtaiRWVdX+1j4ArBpHESdLuJ8UkrwO+Arwoar62bjrGaYk7wIOVtV9465lRFYAFwA3VdX5wH8ypl/nR6WtPW9g7gfZbwOnJ/nT8VY1WjX3rPlYnjc/WcK9+69ESHIqc8F+a1V9ddz1jMDFwLuTPMXcstolSf5xvCUN1V5gb1Ud/o3rDubCvid/APx7Vc1W1X8DXwV+d8w1jcIzSc4BaNuD4yjiZAn3rr8SIUmYW6t9rKo+Pe56RqGqPl5Va6pqkrl/v7urqpu7vqo6AOxJ8sbWtR54dIwljcKPgIuSvLb9n11PZ380bnYCG1t7I7BjHEWM7Vshl9MJ8JUIS3Ux8D7gB0kebH2fqKo7x1eSjsMHgVvbDciTwLVjrmeoqmp3kjuA+5l7wusBXiEf1T9eSb4EvBM4O8le4JPADcDtSTYBTwNXj6U2v35AkvpzsizLSNJJxXCXpA4Z7pLUIcNdkjpkuEtShwx3SeqQ4S5JHfpfchr5WWe14doAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "t = TSTensor(torch.arange(11).float())\n",
    "tt_ = []\n",
    "for _ in range(1000): \n",
    "    tt = TSShuffleSteps()(t, split_idx=0)\n",
    "    test_eq(len(set(tt.tolist())), len(t))\n",
    "    test_ne(tt, t)\n",
    "    tt_.extend([t for i,t in enumerate(tt) if t!=i])\n",
    "x, y = np.unique(tt_, return_counts=True) # This is to visualize distribution which should be equal for all and half for first and last items\n",
    "plt.bar(x, y);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSMagAddNoise(RandTransform):\n",
    "    \"Applies additive noise on the y-axis for each step of a `TSTensor` batch\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=1, ex=None, **kwargs):\n",
    "        self.magnitude, self.ex = magnitude, ex\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        # output = o + torch.normal(0, o.std() * self.magnitude, o.shape, dtype=o.dtype, device=o.device)\n",
    "        output = o + torch.normal(0, 1/3, o.shape, dtype=o.dtype, device=o.device) * (o[..., 1:] - o[..., :-1]).std(2, keepdims=True) * self.magnitude\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output\n",
    "\n",
    "class TSMagMulNoise(RandTransform): \n",
    "    \"Applies multiplicative noise on the y-axis for each step of a `TSTensor` batch\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=1, ex=None, **kwargs): \n",
    "        self.magnitude, self.ex = magnitude, ex\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        noise = torch.normal(1, self.magnitude * .025, o.shape, dtype=o.dtype, device=o.device)\n",
    "        output = o * noise\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSMagAddNoise()(xb, split_idx=0).shape, xb.shape)\n",
    "test_eq(TSMagMulNoise()(xb, split_idx=0).shape, xb.shape)\n",
    "test_ne(TSMagAddNoise()(xb, split_idx=0), xb)\n",
    "test_ne(TSMagMulNoise()(xb, split_idx=0), xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def random_curve_generator(o, magnitude=0.1, order=4, noise=None):\n",
    "    seq_len = o.shape[-1]\n",
    "    f = CubicSpline(np.linspace(-seq_len, 2 * seq_len - 1, 3 * (order - 1) + 1, dtype=int), \n",
    "                    np.random.normal(loc=1.0, scale=magnitude, size=3 * (order - 1) + 1), axis=-1)\n",
    "    return f(np.arange(seq_len))\n",
    "\n",
    "def random_cum_curve_generator(o, magnitude=0.1, order=4, noise=None):\n",
    "    x = random_curve_generator(o, magnitude=magnitude, order=order, noise=noise).cumsum()\n",
    "    x -= x[0]\n",
    "    x /= x[-1]\n",
    "    x = np.clip(x, 0, 1)\n",
    "    return x * (o.shape[-1] - 1)\n",
    "\n",
    "def random_cum_noise_generator(o, magnitude=0.1, noise=None):\n",
    "    seq_len = o.shape[-1]\n",
    "    x = np.clip(np.ones(seq_len) + np.random.normal(loc=0, scale=magnitude, size=seq_len), 0, 1000).cumsum()\n",
    "    x -= x[0]\n",
    "    x /= x[-1]\n",
    "    return x * (o.shape[-1] - 1)\n",
    "\n",
    "def random_cum_linear_generator(o, magnitude=0.1):\n",
    "    seq_len = o.shape[-1]\n",
    "    win_len = int(round(seq_len * np.random.rand() * magnitude))\n",
    "    if win_len == seq_len: return np.arange(o.shape[-1])\n",
    "    start = np.random.randint(0, seq_len - win_len)\n",
    "    # mult between .5 and 2\n",
    "    rand = np.random.rand()\n",
    "    mult = 1 + rand\n",
    "    if np.random.randint(2): mult = 1 - rand/2\n",
    "    x = np.ones(seq_len)\n",
    "    x[start : start + win_len] = mult\n",
    "    x = x.cumsum()\n",
    "    x -= x[0]\n",
    "    x /= x[-1]\n",
    "    return np.clip(x, 0, 1) * (seq_len - 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSTimeNoise(RandTransform):\n",
    "    \"Applies noise to each step in the x-axis of a `TSTensor` batch based on smooth random curve\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.1, ex=None, **kwargs): \n",
    "        self.magnitude, self.ex = magnitude, ex\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        f = CubicSpline(np.arange(o.shape[-1]), o.cpu(), axis=-1)\n",
    "        output = o.new(f(random_cum_noise_generator(o, magnitude=self.magnitude)))\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSTimeNoise()(xb, split_idx=0).shape, xb.shape)\n",
    "test_ne(TSTimeNoise()(xb, split_idx=0), xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSMagWarp(RandTransform):\n",
    "    \"Applies warping to the y-axis of a `TSTensor` batch based on a smooth random curve\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.02, ord=4, ex=None, **kwargs): \n",
    "        self.magnitude, self.ord, self.ex = magnitude, ord, ex\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if self.magnitude and self.magnitude <= 0: return o\n",
    "        y_mult = random_curve_generator(o, magnitude=self.magnitude, order=self.ord)\n",
    "        output = o * o.new(y_mult)\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSMagWarp()(xb, split_idx=0).shape, xb.shape)\n",
    "test_ne(TSMagWarp()(xb, split_idx=0), xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSTimeWarp(RandTransform):\n",
    "    \"Applies time warping to the x-axis of a `TSTensor` batch based on a smooth random curve\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.02, ord=4, ex=None, **kwargs): \n",
    "        self.magnitude, self.ord, self.ex = magnitude, ord, ex\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        f = CubicSpline(np.arange(o.shape[-1]), o.cpu(), axis=-1)\n",
    "        output = o.new(f(random_cum_curve_generator(o, magnitude=self.magnitude, order=self.ord)))\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSTimeWarp()(xb, split_idx=0).shape, xb.shape)\n",
    "test_ne(TSTimeWarp()(xb, split_idx=0), xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSWindowWarp(RandTransform):\n",
    "    \"\"\"Applies window slicing to the x-axis of a `TSTensor` batch based on a random linear curve based on\n",
    "    https://halshs.archives-ouvertes.fr/halshs-01357973/document\"\"\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.1, ex=None, **kwargs): \n",
    "        self.magnitude, self.ex = magnitude, ex\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0 or self.magnitude >= 1: return o\n",
    "        f = CubicSpline(np.arange(o.shape[-1]), o.cpu(), axis=-1)\n",
    "        output = o.new(f(random_cum_linear_generator(o, magnitude=self.magnitude)))\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSWindowWarp()(xb, split_idx=0).shape, xb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSMagScale(RandTransform):\n",
    "    \"Applies scaling to the y-axis of a `TSTensor` batch based on a scalar\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.5, ex=None, **kwargs): \n",
    "        self.magnitude, self.ex = magnitude, ex\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        rand = random_half_normal()\n",
    "        scale = (1 - (rand  * self.magnitude)/2) if random.random() > 1/3 else (1 + (rand  * self.magnitude))\n",
    "        output = o * scale\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output\n",
    "    \n",
    "class TSMagScalePerVar(RandTransform):\n",
    "    \"Applies per_var scaling to the y-axis of a `TSTensor` batch based on a scalar\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.5, ex=None, **kwargs): \n",
    "        self.magnitude, self.ex = magnitude, ex\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        s = [1] * o.ndim\n",
    "        s[-2] = o.shape[-2]\n",
    "        rand = random_half_normal_tensor(s, device=o.device)\n",
    "        scale = (1 - (rand  * self.magnitude)/2) if random.random() > 1/3 else (1 + (rand  * self.magnitude))\n",
    "        output = o * scale\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output\n",
    "    \n",
    "TSMagScaleByVar = TSMagScalePerVar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSMagScale()(xb, split_idx=0).shape, xb.shape)\n",
    "test_eq(TSMagScalePerVar()(xb, split_idx=0).shape, xb.shape)\n",
    "test_ne(TSMagScale()(xb, split_idx=0), xb)\n",
    "test_ne(TSMagScalePerVar()(xb, split_idx=0), xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSRandomResizedCrop(RandTransform):\n",
    "    \"Randomly amplifies a sequence focusing on a random section of the steps\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.1, size=None, scale=None, ex=None, mode='linear', **kwargs): \n",
    "        \"\"\"\n",
    "        Args:\n",
    "            size: None, int or float\n",
    "            scale: None or tuple of 2 floats 0 < float <= 1\n",
    "            mode:  'nearest' | 'linear' | 'area'\n",
    "        \n",
    "        \"\"\"\n",
    "        self.magnitude, self.ex, self.mode = magnitude, ex, mode\n",
    "        if scale is not None: \n",
    "            assert is_listy(scale) and len(scale) == 2 and min(scale) > 0 and min(scale) <= 1, \"scale must be a tuple with 2 floats 0 < float <= 1\"\n",
    "        self.size,self.scale = size,scale\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        seq_len = o.shape[-1]\n",
    "        if self.size is not None: \n",
    "            size = self.size if isinstance(self.size, Integral) else int(round(self.size * seq_len))\n",
    "        else:\n",
    "            size = seq_len\n",
    "        if self.scale is not None: \n",
    "            lambd = np.random.uniform(self.scale[0], self.scale[1])\n",
    "        else: \n",
    "            lambd = np.random.beta(self.magnitude, self.magnitude)\n",
    "            lambd = max(lambd, 1 - lambd)\n",
    "        win_len = int(round(seq_len * lambd))\n",
    "        if win_len == seq_len: \n",
    "            if size == seq_len: return o\n",
    "            _slice = slice(None) \n",
    "        else:\n",
    "            start = np.random.randint(0, seq_len - win_len)\n",
    "            _slice = slice(start, start + win_len)\n",
    "        return F.interpolate(o[..., _slice], size=size, mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)\n",
    "    \n",
    "TSRandomZoomIn = TSRandomResizedCrop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSRandomResizedCrop(.5)(xb, split_idx=0).shape, xb.shape)\n",
    "test_ne(TSRandomResizedCrop(size=.8, scale=(.5, 1))(xb, split_idx=0).shape, xb.shape)\n",
    "test_ne(TSRandomResizedCrop(size=20, scale=(.5, 1))(xb, split_idx=0).shape, xb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSWindowSlicing(RandTransform):\n",
    "    \"Randomly extracts an resize a ts slice based on https://halshs.archives-ouvertes.fr/halshs-01357973/document\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.1, ex=None, mode='linear', **kwargs): \n",
    "        \"mode:  'nearest' | 'linear' | 'area'\"\n",
    "        self.magnitude, self.ex, self.mode = magnitude, ex, mode\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0 or self.magnitude >= 1: return o\n",
    "        seq_len = o.shape[-1]\n",
    "        win_len = int(round(seq_len * (1 - self.magnitude)))\n",
    "        if win_len == seq_len: return o\n",
    "        start = np.random.randint(0, seq_len - win_len)\n",
    "        return F.interpolate(o[..., start : start + win_len], size=seq_len, mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSWindowSlicing()(xb, split_idx=0).shape, xb.shape)\n",
    "test_ne(TSWindowSlicing()(xb, split_idx=0), xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSRandomZoomOut(RandTransform):\n",
    "    \"Randomly compresses a sequence on the x-axis\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.1, ex=None, mode='linear', **kwargs): \n",
    "        \"mode:  'nearest' | 'linear' | 'area'\"\n",
    "        self.magnitude, self.ex, self.mode = magnitude, ex, mode\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        seq_len = o.shape[-1]\n",
    "        lambd = np.random.beta(self.magnitude, self.magnitude)\n",
    "        lambd = max(lambd, 1 - lambd)\n",
    "        win_len = int(round(seq_len * lambd))\n",
    "        if win_len == seq_len: return o\n",
    "        start = (seq_len - win_len) // 2\n",
    "        output = torch.zeros_like(o, dtype=o.dtype, device=o.device)\n",
    "        interp = F.interpolate(o, size=win_len, mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)\n",
    "        output[..., start:start + win_len] = o.new(interp)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSRandomZoomOut(.5)(xb, split_idx=0).shape, xb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSRandomTimeScale(RandTransform):\n",
    "    \"Randomly amplifies/ compresses a sequence on the x-axis keeping the same length\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.1, ex=None, mode='linear', **kwargs): \n",
    "        \"mode:  'nearest' | 'linear' | 'area'\"\n",
    "        self.magnitude, self.ex, self.mode = magnitude, ex, mode\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        if np.random.rand() <= 0.5: return TSRandomZoomIn(magnitude=self.magnitude, ex=self.ex, mode=self.mode)(o, split_idx=0)\n",
    "        else: return TSRandomZoomOut(magnitude=self.magnitude, ex=self.ex, mode=self.mode)(o, split_idx=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSRandomTimeScale(.5)(xb, split_idx=0).shape, xb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSRandomTimeStep(RandTransform):\n",
    "    \"Compresses a sequence on the x-axis by randomly selecting sequence steps and interpolating to previous size\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.02, ex=None, mode='linear', **kwargs): \n",
    "        \"mode:  'nearest' | 'linear' | 'area'\"\n",
    "        self.magnitude, self.ex, self.mode = magnitude, ex, mode\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        seq_len = o.shape[-1]\n",
    "        new_seq_len = int(round(seq_len * max(.5, (1 - np.random.rand() * self.magnitude))))\n",
    "        if  new_seq_len == seq_len: return o\n",
    "        timesteps = np.sort(np.random.choice(np.arange(seq_len),new_seq_len, replace=False))\n",
    "        output = F.interpolate(o[..., timesteps], size=seq_len, mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSRandomTimeStep()(xb, split_idx=0).shape, xb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "\n",
    "class TSResampleSteps(RandTransform):\n",
    "    \"Transform that randomly selects and sorts sequence steps (with replacement) maintaining the sequence length\"\n",
    "\n",
    "    order = 90\n",
    "    def __init__(self, step_pct=1., same_seq_len=True, magnitude=None, **kwargs):\n",
    "        assert step_pct > 0, 'seq_len_pct must be subsample > 0'\n",
    "        self.step_pct, self.same_seq_len = step_pct, same_seq_len\n",
    "        super().__init__(**kwargs)\n",
    "\n",
    "    def encodes(self, o: TSTensor):\n",
    "        S = o.shape[-1]\n",
    "        if isinstance(self.step_pct, tuple):\n",
    "            step_pct = np.random.rand() * (self.step_pct[1] - self.step_pct[0]) + self.step_pct[0]\n",
    "        else: \n",
    "            step_pct = self.step_pct\n",
    "        if step_pct != 1 and self.same_seq_len:\n",
    "            idxs = np.sort(np.tile(np.random.choice(S, round(S * step_pct), True), math.ceil(1 / step_pct))[:S])\n",
    "        else:\n",
    "            idxs = np.sort(np.random.choice(S, round(S * step_pct), True))\n",
    "        return o[..., idxs]\n",
    "    \n",
    "TSSubsampleSteps = TSResampleSteps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSResampleSteps(step_pct=.9, same_seq_len=False)(xb, split_idx=0).shape[-1], round(.9*xb.shape[-1]))\n",
    "test_eq(TSResampleSteps(step_pct=.9, same_seq_len=True)(xb, split_idx=0).shape[-1], xb.shape[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSBlur(RandTransform):\n",
    "    \"Blurs a sequence applying a filter of type [1, 0, 1]\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=1., ex=None, filt_len=None, **kwargs): \n",
    "        self.magnitude, self.ex = magnitude, ex\n",
    "        if filt_len is None: \n",
    "            filterargs = [1, 0, 1]  \n",
    "        else: \n",
    "            filterargs = ([1] * max(1, filt_len // 2) + [0] + [1] * max(1, filt_len // 2))\n",
    "        self.filterargs = np.array(filterargs) \n",
    "        self.filterargs = self.filterargs/self.filterargs.sum()\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        output = o.new(convolve1d(o.cpu(), self.filterargs, mode='nearest'))\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSBlur(filt_len=7)(xb, split_idx=0).shape, xb.shape)\n",
    "test_ne(TSBlur()(xb, split_idx=0), xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSSmooth(RandTransform):\n",
    "    \"Smoothens a sequence applying a filter of type [1, 5, 1]\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=1., ex=None, filt_len=None, **kwargs): \n",
    "        self.magnitude, self.ex = magnitude, ex\n",
    "        self.filterargs = np.array([1, 5, 1])\n",
    "        if filt_len is None: \n",
    "            filterargs = [1, 5, 1]  \n",
    "        else: \n",
    "            filterargs = ([1] * max(1, filt_len // 2) + [5] + [1] * max(1, filt_len // 2))\n",
    "        self.filterargs = np.array(filterargs) \n",
    "        self.filterargs = self.filterargs/self.filterargs.sum()\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        output = o.new(convolve1d(o.cpu(), self.filterargs, mode='nearest'))\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSSmooth(filt_len=7)(xb, split_idx=0).shape, xb.shape)\n",
    "test_ne(TSSmooth()(xb, split_idx=0), xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def maddest(d, axis=None): #Mean Absolute Deviation\n",
    "    return np.mean(np.absolute(d - np.mean(d, axis)), axis)\n",
    "\n",
    "class TSFreqDenoise(RandTransform):\n",
    "    \"Denoises a sequence applying a wavelet decomposition method\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.1, ex=None, wavelet='db4', level=2, thr=None, thr_mode='hard', pad_mode='per', **kwargs): \n",
    "        self.magnitude, self.ex = magnitude, ex\n",
    "        self.wavelet, self.level, self.thr, self.thr_mode, self.pad_mode = wavelet, level, thr, thr_mode, pad_mode\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        \"\"\"\n",
    "        1. Adapted from waveletSmooth function found here:\n",
    "        http://connor-johnson.com/2016/01/24/using-pywavelets-to-remove-high-frequency-noise/\n",
    "        2. Threshold equation and using hard mode in threshold as mentioned\n",
    "        in section '3.2 denoising based on optimized singular values' from paper by Tomas Vantuch:\n",
    "        http://dspace.vsb.cz/bitstream/handle/10084/133114/VAN431_FEI_P1807_1801V001_2018.pdf\n",
    "        \"\"\"\n",
    "        seq_len = o.shape[-1]\n",
    "        # Decompose to get the wavelet coefficients\n",
    "        coeff = pywt.wavedec(o.cpu(), self.wavelet, mode=self.pad_mode)\n",
    "        # Calculate sigma for threshold as defined in http://dspace.vsb.cz/bitstream/handle/10084/133114/VAN431_FEI_P1807_1801V001_2018.pdf\n",
    "        # As noted by @harshit92 MAD referred to in the paper is Mean Absolute Deviation not Median Absolute Deviation\n",
    "        sigma = (1/0.6745) * maddest(coeff[-self.level])\n",
    "        # Calculate the univeral threshold\n",
    "        uthr = sigma * np.sqrt(2*np.log(seq_len)) * (1 if self.thr is None else self.magnitude)\n",
    "        coeff[1:] = (pywt.threshold(c, value=uthr, mode=self.thr_mode) for c in coeff[1:])\n",
    "        # Reconstruct the signal using the thresholded coefficients\n",
    "        output = o.new(pywt.waverec(coeff, self.wavelet, mode=self.pad_mode)[..., :seq_len])\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSFreqDenoise()(xb, split_idx=0).shape, xb.shape)\n",
    "test_ne(TSFreqDenoise()(xb, split_idx=0), xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSRandomFreqNoise(RandTransform):\n",
    "    \"Applys random noise using a wavelet decomposition method\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.1, ex=None, wavelet='db4', level=2, mode='constant', **kwargs): \n",
    "        self.magnitude, self.ex = magnitude, ex\n",
    "        self.wavelet, self.level, self.mode = wavelet, level, mode\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        self.level = 1 if self.level is None else self.level\n",
    "        coeff = pywt.wavedec(o.cpu(), self.wavelet, mode=self.mode, level=self.level)\n",
    "        coeff[1:] = [c * (1 + 2 * (np.random.rand() - 0.5) * self.magnitude) for c in coeff[1:]]\n",
    "        output = o.new(pywt.waverec(coeff, self.wavelet, mode=self.mode)[..., :o.shape[-1]])\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSRandomFreqNoise()(xb, split_idx=0).shape, xb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSRandomResizedLookBack(RandTransform):\n",
    "    \"Selects a random number of sequence steps starting from the end and return an output of the same shape\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.1, mode='linear', **kwargs): \n",
    "        \"mode:  'nearest' | 'linear' | 'area'\"\n",
    "        self.magnitude, self.mode = magnitude, mode\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        seq_len = o.shape[-1]\n",
    "        lambd = np.random.beta(self.magnitude, self.magnitude)\n",
    "        lambd = min(lambd, 1 - lambd)\n",
    "        output = o.clone()[..., int(round(lambd * seq_len)):]\n",
    "        return F.interpolate(output, size=seq_len, mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(100): \n",
    "    o = TSRandomResizedLookBack()(xb, split_idx=0)\n",
    "    test_eq(o.shape[-1], xb.shape[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSRandomLookBackOut(RandTransform):\n",
    "    \"Selects a random number of sequence steps starting from the end and set them to zero\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.1, **kwargs): \n",
    "        self.magnitude = magnitude\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        seq_len = o.shape[-1]\n",
    "        lambd = np.random.beta(self.magnitude, self.magnitude)\n",
    "        lambd = min(lambd, 1 - lambd)\n",
    "        output = o.clone()\n",
    "        output[..., :int(round(lambd * seq_len))] = 0 \n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(100): \n",
    "    o = TSRandomLookBackOut()(xb, split_idx=0)\n",
    "    test_eq(o.shape[-1], xb.shape[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSVarOut(RandTransform):\n",
    "    \"Set the value of a random number of variables to zero\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.05, ex=None, **kwargs): \n",
    "        self.magnitude, self.ex = magnitude, ex\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        in_vars = o.shape[-2]\n",
    "        if in_vars == 1: return o\n",
    "        lambd = np.random.beta(self.magnitude, self.magnitude)\n",
    "        lambd = min(lambd, 1 - lambd)\n",
    "        p = np.arange(in_vars).cumsum()\n",
    "        p = p/p[-1]\n",
    "        p = p / p.sum()\n",
    "        p = p[::-1]\n",
    "        out_vars = np.random.choice(np.arange(in_vars), int(round(lambd * in_vars)), p=p, replace=False)\n",
    "        if len(out_vars) == 0:  return o\n",
    "        output = o.clone()\n",
    "        output[...,out_vars,:] = 0\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSVarOut()(xb, split_idx=0).shape, xb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSCutOut(RandTransform):\n",
    "    \"Sets a random section of the sequence to zero\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.05, ex=None, **kwargs): \n",
    "        self.magnitude, self.ex = magnitude, ex\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        seq_len = o.shape[-1]\n",
    "        lambd = np.random.beta(self.magnitude, self.magnitude)\n",
    "        lambd = min(lambd, 1 - lambd)\n",
    "        win_len = int(round(seq_len * lambd))\n",
    "        start = np.random.randint(-win_len + 1, seq_len)\n",
    "        end = start + win_len\n",
    "        start = max(0, start)\n",
    "        end = min(end, seq_len)\n",
    "        output = o.clone()\n",
    "        output[..., start:end] = 0\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSCutOut()(xb, split_idx=0).shape, xb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSTimeStepOut(RandTransform):\n",
    "    \"Sets random sequence steps to zero\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.05, ex=None, **kwargs): \n",
    "        self.magnitude, self.ex = magnitude, ex\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        magnitude = min(.5, self.magnitude)\n",
    "        seq_len = o.shape[-1]\n",
    "        timesteps = np.sort(np.random.choice(np.arange(seq_len), int(round(seq_len * magnitude)), replace=False))\n",
    "        output = o.clone()\n",
    "        output[..., timesteps] = 0\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSTimeStepOut()(xb, split_idx=0).shape, xb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSRandomCropPad(RandTransform):\n",
    "    \"Crops a section of the sequence of a random length\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.05, ex=None, **kwargs): \n",
    "        self.magnitude, self.ex = magnitude, ex\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        seq_len = o.shape[-1]\n",
    "        lambd = np.random.beta(self.magnitude, self.magnitude)\n",
    "        lambd = max(lambd, 1 - lambd)\n",
    "        win_len = int(round(seq_len * lambd))\n",
    "        if win_len == seq_len: return o\n",
    "        start = np.random.randint(0, seq_len - win_len)\n",
    "        output = torch.zeros_like(o, dtype=o.dtype, device=o.device)\n",
    "        output[..., start : start + win_len] = o[..., start : start + win_len]\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSRandomCropPad()(xb, split_idx=0).shape, xb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "\n",
    "from tsai.callback.MVP import create_mask\n",
    "\n",
    "class TSMaskOut(RandTransform):\n",
    "    \"\"\"Applies a random mask\"\"\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.1, lm:int=3, stateful:bool=True, sync:bool=False, subsequence_mask:bool=True, \n",
    "                 variable_mask:bool=False, future_mask:bool=False, schedule_func:Optional[callable]=None, compensate:bool=False, ex=None, **kwargs):\n",
    "        store_attr()\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        r = self.magnitude * self.schedule_func(self.pct_train) if self.schedule_func is not None else self.magnitude\n",
    "        mask = create_mask(o,  r=r, lm=self.lm, stateful=self.stateful, sync=self.sync, \n",
    "                           subsequence_mask=self.subsequence_mask, variable_mask=self.variable_mask, future_mask=self.future_mask)\n",
    "        if self.compensate: # per sample and feature\n",
    "            mean_per_seq = (torch.max(torch.ones(1, device=mask.device), torch.sum(mask, dim=-1).unsqueeze(-1)) / mask.shape[-1])\n",
    "            output = o.masked_fill(mask, 0) / (1 - mean_per_seq)\n",
    "        else:\n",
    "            output = o.masked_fill(mask, 0)\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/nacho/opt/anaconda3/envs/py37/lib/python3.7/site-packages/torch/_tensor.py:579: UserWarning: floor_divide is deprecated, and will be removed in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values.\n",
      "To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor'). (Triggered internally at  ../aten/src/ATen/native/BinaryOps.cpp:467.)\n",
      "  return torch.floor_divide(other, self)\n"
     ]
    }
   ],
   "source": [
    "test_eq(TSMaskOut()(xb, split_idx=0).shape, xb.shape)\n",
    "test_ne(TSMaskOut()(xb, split_idx=0), xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "\n",
    "class TSInputDropout(RandTransform):\n",
    "    \"\"\"Applies input dropout with required_grad=False\"\"\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0., ex=None, **kwargs):\n",
    "        self.magnitude, self.ex = magnitude, ex\n",
    "        self.dropout = nn.Dropout(magnitude)\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        with torch.no_grad(): output = self.dropout(o)\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSInputDropout(.1)(xb, split_idx=0).shape, xb.shape)\n",
    "test_ne(TSInputDropout(.1)(xb, split_idx=0), xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSTranslateX(RandTransform):\n",
    "    \"Moves a selected sequence window a random number of steps\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.1, ex=None, **kwargs): \n",
    "        self.magnitude, self.ex = magnitude, ex\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        seq_len = o.shape[-1]\n",
    "        lambd = np.random.beta(self.magnitude, self.magnitude)\n",
    "        lambd = min(lambd, 1 - lambd)\n",
    "        shift = int(round(seq_len * lambd))\n",
    "        if shift == 0 or shift == seq_len: return o\n",
    "        if np.random.rand() < 0.5: shift = -shift\n",
    "        new_start = max(0, shift)\n",
    "        new_end = min(seq_len + shift, seq_len)\n",
    "        start = max(0, -shift)\n",
    "        end = min(seq_len - shift, seq_len)\n",
    "        output = torch.zeros_like(o, dtype=o.dtype, device=o.device)\n",
    "        output[..., new_start : new_end] = o[..., start : end]\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSTranslateX()(xb, split_idx=0).shape, xb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSRandomShift(RandTransform):\n",
    "    \"Shifts and splits a sequence\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.02, ex=None, **kwargs): \n",
    "        self.magnitude, self.ex = magnitude, ex\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        pos = int(round(np.random.randint(0, o.shape[-1]) * self.magnitude)) * (random.randint(0, 1)*2-1)\n",
    "        output = torch.cat((o[..., pos:], o[..., :pos]), dim=-1)\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSRandomShift()(xb, split_idx=0).shape, xb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSHorizontalFlip(RandTransform):\n",
    "    \"Flips the sequence along the x-axis\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=1., ex=None, **kwargs): \n",
    "        self.magnitude, self.ex = magnitude, ex\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        output = torch.flip(o, [-1])\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSHorizontalFlip()(xb, split_idx=0).shape, xb.shape)\n",
    "test_ne(TSHorizontalFlip()(xb, split_idx=0), xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSRandomTrend(RandTransform):\n",
    "    \"Randomly rotates the sequence along the z-axis\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.1, ex=None, **kwargs): \n",
    "        self.magnitude, self.ex = magnitude, ex\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        flat_x = o.reshape(o.shape[0], -1)\n",
    "        ran = flat_x.max(dim=-1, keepdim=True)[0] - flat_x.min(dim=-1, keepdim=True)[0]\n",
    "        trend = torch.linspace(0, 1, o.shape[-1], device=o.device) * ran\n",
    "        t = (1 + self.magnitude * 2 * (np.random.rand() - 0.5) * trend)\n",
    "        t -= t.mean(-1, keepdim=True)\n",
    "        if o.ndim == 3: t = t.unsqueeze(1)\n",
    "        output = o + t\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output\n",
    "    \n",
    "TSRandomRotate = TSRandomTrend"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSRandomTrend()(xb, split_idx=0).shape, xb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSVerticalFlip(RandTransform):\n",
    "    \"Applies a negative value to the time sequence\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=1., ex=None, **kwargs): \n",
    "        self.magnitude, self.ex = magnitude, ex\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor): \n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        return - o"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSVerticalFlip()(xb, split_idx=0).shape, xb.shape)\n",
    "test_ne(TSVerticalFlip()(xb, split_idx=0), xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSResize(RandTransform):\n",
    "    \"Resizes the sequence length of a time series\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=-0.5, size=None, ex=None, mode='linear', **kwargs): \n",
    "        \"mode:  'nearest' | 'linear' | 'area'\"\n",
    "        self.magnitude, self.size, self.ex, self.mode = magnitude, size, ex, mode\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor): \n",
    "        if self.magnitude == 0: return o\n",
    "        size = ifnone(self.size, int(round((1 + self.magnitude) * o.shape[-1])))\n",
    "        output = F.interpolate(o, size=size, mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for sz in np.linspace(.2, 2, 10): test_eq(TSResize(sz)(xb, split_idx=0).shape[-1], int(round(xb.shape[-1]*(1+sz))))\n",
    "test_ne(TSResize(1)(xb, split_idx=0), xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSRandomSize(RandTransform):\n",
    "    \"Randomly resizes the sequence length of a time series\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.1, ex=None, mode='linear', **kwargs):\n",
    "        \"mode:  'nearest' | 'linear' | 'area'\"\n",
    "        self.magnitude, self.ex, self.mode = magnitude, ex, mode\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        size_perc = 1 + random_half_normal() * self.magnitude * (-1 if random.random() > .5 else 1)\n",
    "        return F.interpolate(o, size=int(size_perc * o.shape[-1]), mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "seq_len_ = []\n",
    "for i in range(100): \n",
    "    o = TSRandomSize(.5)(xb, split_idx=0)\n",
    "    seq_len_.append(o.shape[-1])\n",
    "test_lt(min(seq_len_), xb.shape[-1])\n",
    "test_gt(max(seq_len_), xb.shape[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSRandomLowRes(RandTransform):\n",
    "    \"Randomly resizes the sequence length of a time series to a lower resolution\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=.5, ex=None, mode='linear', **kwargs): \n",
    "        \"mode:  'nearest' | 'linear' | 'area'\"\n",
    "        self.magnitude, self.ex, self.mode = magnitude, ex, mode\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor): \n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        size_perc = 1 - (np.random.rand() * (1 - self.magnitude))\n",
    "        return F.interpolate(o, size=int(size_perc * o.shape[-1]), mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSDownUpScale(RandTransform):\n",
    "    \"Downscales a time series and upscales it again to previous sequence length\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.5, ex=None, mode='linear', **kwargs): \n",
    "        \"mode:  'nearest' | 'linear' | 'area'\"\n",
    "        self.magnitude, self.ex, self.mode = magnitude, ex, mode\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor): \n",
    "        if not self.magnitude or self.magnitude <= 0 or self.magnitude >= 1: return o\n",
    "        output = F.interpolate(o, size=int((1 - self.magnitude) * o.shape[-1]), mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)\n",
    "        output = F.interpolate(output, size=o.shape[-1], mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSDownUpScale()(xb, split_idx=0).shape, xb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TSRandomDownUpScale(RandTransform):\n",
    "    \"Randomly downscales a time series and upscales it again to previous sequence length\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=.5, ex=None, mode='linear', **kwargs): \n",
    "        \"mode:  'nearest' | 'linear' | 'area'\"\n",
    "        self.magnitude, self.ex, self.mode = magnitude, ex, mode\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor): \n",
    "        if not self.magnitude or self.magnitude <= 0 or self.magnitude >= 1: return o\n",
    "        scale_factor = 0.5 + 0.5 * np.random.rand() \n",
    "        output = F.interpolate(o, size=int(scale_factor * o.shape[-1]), mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)\n",
    "        output = F.interpolate(output, size=o.shape[-1], mode=self.mode, align_corners=None if self.mode in ['nearest', 'area'] else False)\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(TSRandomDownUpScale()(xb, split_idx=0).shape, xb.shape)\n",
    "test_ne(TSDownUpScale()(xb, split_idx=0), xb)\n",
    "test_eq(TSDownUpScale()(xb, split_idx=1), xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "\n",
    "class TSRandomConv(RandTransform):\n",
    "    \"\"\"Applies a convolution with a random kernel and random weights with required_grad=False\"\"\"\n",
    "    order = 90\n",
    "    def __init__(self, magnitude=0.05, ex=None, ks=[1, 3, 5, 7], **kwargs):\n",
    "        self.magnitude, self.ex, self.ks = magnitude, ex, ks\n",
    "        self.conv = nn.Conv1d(1, 1, 1, bias=False)\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor):\n",
    "        if not self.magnitude or self.magnitude <= 0 or self.ks is None: return o\n",
    "        ks = np.random.choice(self.ks, 1)[0] if is_listy(self.ks) else self.ks\n",
    "        c_in = o.shape[1]\n",
    "        weight = nn.Parameter(torch.zeros(c_in, c_in, ks, device=o.device, requires_grad=False))\n",
    "        nn.init.kaiming_normal_(weight)\n",
    "        self.conv.weight = weight\n",
    "        self.conv.padding = ks // 2\n",
    "        output = (1 - self.magnitude) * o + self.magnitude * self.conv(o)\n",
    "        if self.ex is not None: output[...,self.ex,:] = o[...,self.ex,:]\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(5): \n",
    "    o = TSRandomConv(magnitude=0.05, ex=None, ks=[1, 3, 5, 7])(xb, split_idx=0)\n",
    "    test_eq(o.shape, xb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "all_TS_randaugs = [\n",
    "    \n",
    "    TSIdentity, \n",
    "    \n",
    "    # Noise\n",
    "    (TSMagAddNoise, 0.1, 1.),\n",
    "    (partial(TSMagMulNoise, ex=0), 0.1, 1),\n",
    "    (partial(TSTimeNoise, ex=0), 0.1, 1.),\n",
    "    (partial(TSRandomFreqNoise, ex=0), 0.1, 1.),\n",
    "    partial(TSShuffleSteps, ex=0),\n",
    "    (TSRandomTimeScale, 0.05, 0.5), \n",
    "    (TSRandomTimeStep, 0.05, 0.5), \n",
    "    (partial(TSFreqDenoise, ex=0), 0.1, 1.),\n",
    "    (TSRandomLowRes, 0.05, 0.5),\n",
    "    (TSInputDropout, 0.05, .5),\n",
    "    \n",
    "    # Magnitude\n",
    "    (partial(TSMagWarp, ex=0), 0.02, 0.2),\n",
    "    (TSMagScale, 0.2, 1.),\n",
    "    (partial(TSMagScalePerVar, ex=0), 0.2, 1.),\n",
    "    (partial(TSRandomConv, ex=0), .05, .2),\n",
    "    partial(TSBlur, ex=0),\n",
    "    partial(TSSmooth, ex=0),\n",
    "    partial(TSDownUpScale, ex=0),\n",
    "    partial(TSRandomDownUpScale, ex=0), \n",
    "    (TSRandomTrend, 0.1, 0.5), \n",
    "    TSVerticalFlip, \n",
    "    (TSVarOut, 0.05, 0.5), \n",
    "    (TSCutOut, 0.05, 0.5), \n",
    "    \n",
    "    # Time\n",
    "    (partial(TSTimeWarp, ex=0), 0.02, 0.2),\n",
    "    (TSWindowWarp, 0.05, 0.5),\n",
    "    (TSRandomSize, 0.05, 1.),\n",
    "    TSHorizontalFlip, \n",
    "    (TSTranslateX, 0.1, 0.5),\n",
    "    (TSRandomShift, 0.02, 0.2), \n",
    "    (TSRandomZoomIn, 0.05, 0.5), \n",
    "    (TSWindowSlicing, 0.05, 0.2),\n",
    "    (TSRandomZoomOut, 0.05, 0.5),\n",
    "    (TSRandomLookBackOut, 0.1, 1.),\n",
    "    (TSRandomResizedLookBack, 0.1, 1.),\n",
    "    (TSTimeStepOut, 0.01, 0.2),\n",
    "    (TSRandomCropPad, 0.05, 0.5), \n",
    "    (TSRandomResizedCrop, 0.05, 0.5),\n",
    "    (TSMaskOut, 0.01, 0.2),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class RandAugment(RandTransform):\n",
    "    order = 90\n",
    "    def __init__(self, tfms:list, N:int=1, M:int=3, **kwargs):\n",
    "        '''\n",
    "        tfms   : list of tfm functions (not called)\n",
    "        N      : number of tfms applied to each batch (usual values 1-3)\n",
    "        M      : tfm magnitude multiplier (1-10, usually 3-5). Only works if tfms are tuples (tfm, min, max)\n",
    "        kwargs : RandTransform kwargs\n",
    "        '''\n",
    "        super().__init__(**kwargs)\n",
    "        if not isinstance(tfms, list): tfms = [tfms]\n",
    "        self.tfms, self.N, self.magnitude = tfms, min(len(tfms), N), M / 10\n",
    "        self.n_tfms, self.tfms_idxs = len(tfms), np.arange(len(tfms))\n",
    "\n",
    "    def encodes(self, o:(NumpyTensor, TSTensor)):\n",
    "        if not self.N or not self.magnitude: return o\n",
    "        tfms = self.tfms if self.n_tfms==1 else L(self.tfms)[np.random.choice(np.arange(self.n_tfms), self.N, replace=False)]\n",
    "        tfms_ = []\n",
    "        for tfm in tfms:\n",
    "            if isinstance(tfm, tuple):\n",
    "                t, min_val, max_val = tfm\n",
    "                tfms_ += [t(magnitude=self.magnitude * float(max_val - min_val) + min_val)]\n",
    "            else:  tfms_ += [tfm()]\n",
    "        output = compose_tfms(o, tfms_, split_idx=self.split_idx)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_ne(RandAugment(TSMagAddNoise, N=5, M=10)(xb, split_idx=0), xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TestTfm(RandTransform):\n",
    "    \"Utility class to test the output of selected tfms during training\"\n",
    "    def __init__(self, tfm, magnitude=1., ex=None, **kwargs): \n",
    "        self.tfm, self.magnitude, self.ex = tfm, magnitude, ex\n",
    "        self.tfmd, self.shape = [], []\n",
    "        super().__init__(**kwargs)\n",
    "    def encodes(self, o: TSTensor): \n",
    "        if not self.magnitude or self.magnitude <= 0: return o\n",
    "        output = self.tfm(o, split_idx=self.split_idx)\n",
    "        self.tfmd.append(torch.equal(o, output))\n",
    "        self.shape.append(o.shape)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def get_tfm_name(tfm):\n",
    "    if isinstance(tfm, tuple): tfm = tfm[0]\n",
    "    if hasattr(tfm, \"func\"): tfm = tfm.func\n",
    "    if hasattr(tfm, \"__name__\"): return tfm.__name__\n",
    "    elif hasattr(tfm, \"__class__\") and hasattr(tfm.__class__, \"__name__\"): return tfm.__class__.__name__\n",
    "    else: return tfm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(get_tfm_name(partial(TSMagScale()))==get_tfm_name((partial(TSMagScale()), 0.1, .05))==get_tfm_name(TSMagScale())==get_tfm_name((TSMagScale(), 0.1, .05)), True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_TS_randaugs_names = [get_tfm_name(t) for t in all_TS_randaugs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src onerror=\"\n",
       "        this.nextElementSibling.focus();\n",
       "        this.dispatchEvent(new KeyboardEvent('keydown', {key:'s', keyCode: 83, metaKey: true}));\n",
       "        \" style=\"display:none\"><input style=\"width:0;height:0;border:0\">"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<img src onerror=\"\n",
       "        this.nextElementSibling.focus();\n",
       "        this.dispatchEvent(new KeyboardEvent('keydown', {key:'s', keyCode: 83, metaKey: true}));\n",
       "        \" style=\"display:none\"><input style=\"width:0;height:0;border:0\">"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converted 017_data.transforms.ipynb.\n",
      "\n",
      "\n",
      "Correct conversion! 😃\n",
      "Total time elapsed 0.647 s\n",
      "Thursday 30/09/21 20:55:29 CEST\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio  controls=\"controls\" autoplay=\"autoplay\">\n",
       "                    <source src=\"data:audio/wav;base64,UklGRvQHAABXQVZFZm10IBAAAAABAAEAECcAACBOAAACABAAZGF0YdAHAAAAAPF/iPh/gOoOon6w6ayCoR2ZeyfbjobxK+F2Hs0XjKc5i3DGvzaTlEaraE+zz5uLUl9f46fHpWJdxVSrnfmw8mYEScqUP70cb0Q8X41uysJ1si6Eh1jYzXp9IE2DzOYsftYRyoCY9dJ/8QICgIcEun8D9PmAaBPlfT7lq4MFIlh61tYPiCswIHX+yBaOqT1QbuW7qpVQSv9lu6+xnvRVSlyopAypbGBTUdSalrSTaUBFYpInwUpxOzhti5TOdndyKhCGrdwAfBUcXIJB69p+Vw1egB76+n9q/h6ADglbf4LvnIHfF/981ODThF4m8HiS0riJVjQ6c+/EOZCYQfJrGrhBmPVNMmNArLKhQlkXWYqhbaxXY8ZNHphLuBJsZUEckCTFVHMgNKGJytIDeSUmw4QN4Qx9pReTgb3vYX/TCBuApf75f+P5Y4CRDdN+B+tngk8c8nt03CKGqipgd13OhotwOC5x9MCAknFFcmlmtPmagFFFYOCo0qRzXMhVi57pryNmIEqJlRi8bm52PfuNM8k4dfQv+4cO12l6zCGdg3jl730uE/KAPvS+f0wEAoAsA89/XfXQgBESIn6S5luDtiC8eh/YmIfpLqt1OMp5jXg8/24MveqUNUnPZsqw0Z3yVDldnaUOqIZfXlKrm36zzWhjRhaT+r+ncHI5/otUzfd2uSt7hl/bqXtoHaCC6+mqfrAOeoDD+PJ/xf8RgLMHfH/b8GeBihZIfSXidoQSJWB52NM1iRkzz3MkxpKPbUCrbDu5d5fgTAxkSK3JoEhYD1p2omere2LZTuqYLbdWa49Cx5Dww7tyXDUnioXRkHhwJyKFvd/AfPoYy4Fl7j1/LQorgEr9/X89+0qAOAwAf13sJoL8Gkd8wt25hWIp3Heez/eKODfPcSPCzpFNRDVqf7UlmnNQKGHgqd+jgVvJVm2f265QZTpLS5byur1tpT6ajvrHq3Q2MXWIxtUCehoj8YMk5LB9hRQegeTypn+nBQWA0QHgf7f2q4C5EFt+5ucOg2YfHXtq2SSHpS0ydnTL4IxFO6pvNb4ulBdInWfcsfSc7VMmXpSmE6eeXmZThJxpsgRohEfOk86+AHCoOpOMFsx1dv8s6oYT2k17uR7ngpXod34IEJqAaPfnfyABCIBZBpl/NPI2gTQVjX134x2ExSPMeR7VtYjZMWJ0W8ftjkA/YW1durCWykvjZFKu4p9LVwVbZKNkqpxh6U+6mRC2mGq2Q3SRvsIgcpc2sIpD0Bp4uiiFhW3ecXxOGgaCDe0Vf4cLPoDv+/5/mfw1gN4KKX+17emBqBmYfBHfVYUZKFR44NBtiv41bHJUwx+RJkP1apu2VJlkTwli4qrwoo1ax1dToNCtemRSTBGXz7kJbdM/PY/Dxht0dTLziH7Ul3loJEiE0uJsfdsVTYGL8Yt/AgcMgHYA7X8S+IqAYA+QfjzpxIIVHnp7tdqzhmAstXaxzEqMETpScGC/dJP3Rmdo8LIZnOVSEF+Opxumsl1sVF+dVrE5Z6NIiZSkvVdv2zsqjdnK8HVDLlyHyNjuegogM4NA5z9+YRG9gA722H97AgOA/gSyf43zCIHdE899yuTIg3ciNXpm1jmImTDwdJPITI4RPhRugbvslbFKt2Vfr/6eTFb4W1WkY6m6YPdQjJr2tNZp3EQlko7BgXHRNz2LAc+gdwMq7IUf3R58ohtFgrbr6n7hDFWAlPr8f/T9I4CECU9/De+vgVQY5nxh4POEzybJeCTS5YnCNAZzhsRzkP1Bsmu4t4aYU07nYuerA6KWWcJYO6HHrKJjaE3Zl624UWz/QOOPjcWHc7QzdIk40yl5tCWjhIDhJX0xF4CBMvBsf10IF4Ac//Z/bPlsgAcOwn6S6n6CwxzUewLcRoYaKzV38M23i9o493CNwL6S1UUuaQe0QpvbUfdfiqglpcRccFU+nkWwambASUiVfLyqbg49xY2eyWh1hy/Sh37XjHpaIYKD7OUEfrgS5IC09MV/1gMBgKMDyH/n9N6AhhINfh7mdoMoIZt6r9fAh1cvfHXNya6N4DzDbqi8K5WWSYlmbbAdnkpV6FxJpWSo1V8DUmGb3rMRaQBG2JJgwN9wCDnNi8HNI3dKK1aG0dvHe/UciIJf6rt+Og5wgDn59X9P/xWAKQhxf2XweYH+FjB9suGVhIMlOnlo02GJhTOdc7vFyo/TQGxs2Li7lz9NwmPurBihnVi7WSWiwKvGYntOpJiOt5drKUKMkFnE8HLxNPmJ9NG4eP8mAYUv4Np8hhi3gdruSX+3CSWAwP38f8f6UoCuDPF+6Os8gnAbKnxQ3d2F0imydzDPKIuiN5lxu8EKkrFE82kftW2az1DbYImpMqTUW3FWIJ83r5hl2koJlla7+m0+PmSOZcjcdMgwS4g11iZ6qCLUg5jkxn0QFA6BWvOvfzEFBIBHAtp/Qfa3gC4RSH5y5yeD2B/8evnYS4cULgR2CMsUja47cG/QvW6UeEhXZ3+xP51GVNVdP6Zpp+1eDFM5nMeySWghR4+TNL85cD46YIyCzKJ2kCzEhoTabXtGHs+CCemJfpMPjoDe9+t/qQALgM8Gj3++8UaBqRV2fQTjO4Q3JKd5r9TgiEYyMHTxxiWPpz8jbfq585YpTJpk960xoKFXsVoTo7yq6GGMTw==\" type=\"audio/wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#hide\n",
    "from tsai.imports import create_scripts\n",
    "from tsai.export import get_nb_name\n",
    "nb_name = get_nb_name()\n",
    "create_scripts(nb_name);"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
